#include "MCUT.h"
#include "../cmn/cmn.h"


MCUT::MCUT(int color_num)
{
	kd_tree.valid_node_num = 0;
	kd_tree.leaf_num       = color_num;
	kd_tree.layer_num      = log2(2 * color_num);
	kd_tree.max_node_num   = 2 * color_num - 1;
	kd_tree.p_node         = new NODE [kd_tree.max_node_num];

	for (int i = 0; i < kd_tree.max_node_num; i++)
	{
		kd_tree.p_node[i].dim       = -1;
		kd_tree.p_node[i].th        = -1.0f;
		kd_tree.p_node[i].pixel_num =  0;

		for (int j = 0; j < MAX_COLOR_CHANNEL; j++)
		{
			kd_tree.p_node[i].size[j]  = 0;
			kd_tree.p_node[i].color[j] = 0;
			kd_tree.p_node[i].start[j] = 0;
			kd_tree.p_node[i].end[j]   = 0;
		}
	}
	for (int i = 0; i < MAX_COLOR_CHANNEL; i++)
	{
		p_hist1d[i] = new unsigned int[MAX_COLOR_DIMENTION];
	}
}


MCUT::~MCUT(void)
{
	delete[] kd_tree.p_node;
}


void MCUT::create_1d_hist_for_color_image()
{
	int x = 0, y = 0;
	unsigned char p_pixel[3];

	for (int ci = 0; ci < MAX_COLOR_CHANNEL; ci++)
	{
		for (int di = 0; di < MAX_COLOR_DIMENTION; di++)
		{
			p_hist1d[ci][di] = 0;
		}
	}

	for (y = 0; y < img_height; y++)
	{
		for (x = 0; x < img_width; x++)
		{
			access_pixel(img_width, p_src_img, 24, p_pixel, x, y);
			for (int ci = 0; ci < MAX_COLOR_CHANNEL; ci++)
			{
				p_hist1d[ci][p_pixel[ci]]++;
			}
		}
	}
}

void MCUT::find_suit_color_for_leaf_nodes_by_src_img()
{
	int x = 0, y = 0;
	unsigned char p_pixel[MAX_COLOR_CHANNEL];
	for (int i = 0; i < kd_tree.max_node_num; i++)
	{
		kd_tree.p_node[i].pixel_num = 0;
		for (int j = 0; j < MAX_COLOR_CHANNEL; j++)
		{
			kd_tree.p_node[i].color[j] = 0;
		}
	}


	for (y = 0; y < img_height; y++)
	{
		for (x = 0; x < img_width; x++)
		{
			access_pixel(img_width, p_src_img, 24, p_pixel, x, y);
			int index = index_kd_tree(p_pixel);
			kd_tree.p_node[index].pixel_num++;
			for (int ci = 0; ci < MAX_COLOR_CHANNEL; ci++)
			{
				kd_tree.p_node[index].color[ci] += p_pixel[ci];
			}
		}
	}

	for (int i = 0; i < kd_tree.max_node_num; i++)
	{
		for (int j = 0; j < MAX_COLOR_CHANNEL && kd_tree.p_node[i].pixel_num > 0; j++)
		{
			kd_tree.p_node[i].color[j] = (int)((float)kd_tree.p_node[i].color[j] / (float)(kd_tree.p_node[i].pixel_num));
		}
	}

}

void MCUT::map_src_img_to_cut_img(unsigned char* p_dst_img)
{
	int x = 0, y = 0;
	unsigned char p_pixel[MAX_COLOR_CHANNEL];
	for (y = 0; y < img_height; y++)
	{
		for (x = 0; x < img_width; x++)
		{
			access_pixel(img_width, p_src_img, 24, p_pixel, x, y);
			int index = index_kd_tree(p_pixel);
			for (int j = 0; j < MAX_COLOR_CHANNEL; j++)
			{
				p_pixel[j] = kd_tree.p_node[index].color[j];
			}
			set_pixel(img_width, p_dst_img, 24, p_pixel, x, y);
		}
	}
}

void MCUT::map_src_img_to_cut_img_for_superpixel(S_Pix* p_spix, int spix_num, short* p_spix_label_map, unsigned char* p_dst_img)
{
	int x = 0, y = 0;
	unsigned char p_pixel[MAX_COLOR_CHANNEL];
	for (y = 0; y < img_height; y++)
	{
		for (x = 0; x < img_width; x++)
		{
			int si = p_spix_label_map[y * img_width + x];
			p_pixel[0] =  p_spix[si].color.c0;
			p_pixel[1] =  p_spix[si].color.c1;
			p_pixel[2] =  p_spix[si].color.c2;
			int index = index_kd_tree(p_pixel);
			for (int j = 0; j < MAX_COLOR_CHANNEL; j++)
			{
				p_pixel[j] = kd_tree.p_node[index].color[j];
			}
			set_pixel(img_width, p_dst_img, 24, p_pixel, x, y);
		}
	}
}



void MCUT::create_1d_hist_for_spix(S_Pix* p_spix, int spix_num)
{
	for (int ci = 0; ci < MAX_COLOR_CHANNEL; ci++)
	{
		for (int di = 0; di < MAX_COLOR_DIMENTION; di++)
		{
			p_hist1d[ci][di] = 0;
		}
	}
	
	for (int si = 0; si < spix_num; si++)
	{
		p_hist1d[0][p_spix[si].color.c0]++;
		p_hist1d[1][p_spix[si].color.c1]++;
		p_hist1d[2][p_spix[si].color.c2]++;
	}
}


void MCUT::copy_node(NODE* p_src, NODE* p_dst)
{
	memcpy((unsigned char*)p_dst, (unsigned char*)p_src, sizeof(NODE));
}

void MCUT::spit_node(NODE* p_parent_node, NODE* p_left_node, NODE* p_right_node)
{
	int result_size = 0;
	copy_node(p_parent_node, p_left_node);
	copy_node(p_parent_node, p_right_node);

	p_left_node ->end[p_parent_node->dim]   = (int) p_parent_node->th; 
	p_right_node->start[p_parent_node->dim] = (int)(p_parent_node->th) + 1; 

	result_size = (int)((float)p_left_node->size[p_parent_node->dim] / 2.0f + 0.5f);

	p_left_node-> size[p_parent_node->dim] = result_size;
	p_right_node->size[p_parent_node->dim] = result_size;
}

void MCUT::get_child_id(int parent_id, int* p_left_child_id, int* p_right_child_id)
{
	(*p_left_child_id)  = 2 * parent_id + 1;
	(*p_right_child_id) = 2 * parent_id + 2;
}

void MCUT::get_max_dimesion(NODE* p_node)
{
	int max_size = -1, max_dim = 0;
	for (int ci = 0; ci < MAX_COLOR_CHANNEL; ci++)
	{
		if (p_node->size[ci] > max_size)
		{
			max_size = p_node->size[ci];
			max_dim  = ci;
		}
	}
	p_node->dim = max_dim;
}


void MCUT::init_root_node(NODE* p_root)
{
	int size[MAX_COLOR_CHANNEL] = {0};
	p_root->dim   = 0;
	for (int ci = 0; ci < MAX_COLOR_CHANNEL; ci++)
	{
		for (int di = 0; di < MAX_COLOR_DIMENTION; di++)
		{
			if (p_hist1d[ci][di] > 0) size[ci]++;
		}
		p_root->size[ci]  = size[ci];
		p_root->start[ci] = 0;
		p_root->end[ci]   = MAX_COLOR_DIMENTION;
	}
}


void MCUT::find_median(NODE* p_node)
{
	int th    = 0;
	int end   = p_node->end[p_node->dim];
	int start = p_node->start[p_node->dim];
	int half_size = (int)((float)p_node->size[p_node->dim] / 2.0f + 0.5f);
	for (int di = start; di < end; di++)
	{
		if (p_hist1d[p_node->dim][di] > 0) 
		{
			th++;
		}
		if (th >= half_size)
		{
			p_node->th = (float)(th + start) + 0.5f;
			break;
		}
	}
}

void MCUT::create_kd_tree()
{
	int parent_id, left_id, right_id;

	init_root_node(&kd_tree.p_node[0]);
	kd_tree.valid_node_num = 0;

	for (int depth = 0; depth < kd_tree.layer_num - 1; depth++)
	{
		int parent_start = pow(2, depth) - 1;
		int parent_end   = pow(2, depth+1) - 1;
		for (int parent_id = parent_start; parent_id < parent_end; parent_id++)
		{
			get_child_id(parent_id, &left_id, &right_id);
			NODE* p_parent_node = &kd_tree.p_node[parent_id];
			NODE* p_left_node   = &kd_tree.p_node[left_id];
			NODE* p_right_node  = &kd_tree.p_node[right_id];
			get_max_dimesion(p_parent_node);
			find_median(p_parent_node);
			spit_node(p_parent_node, p_left_node, p_right_node);
			kd_tree.valid_node_num = right_id + 1;

			//Error Test
			if (p_left_node ->size[p_parent_node->dim] < 2 ||
			    p_right_node->size[p_parent_node->dim] < 2 ||
				kd_tree.valid_node_num > kd_tree.max_node_num)
			{
				depth     = kd_tree.layer_num;
				parent_id = parent_end;
			}

		}
	}
}

int MCUT::index_kd_tree(unsigned char* p_pixel)
{
	int parent_id, left_id, right_id;
	int index = 0;
	parent_id = 0;
	NODE* p_node = &kd_tree.p_node[0];
	for (int depth = 0; depth < kd_tree.layer_num - 1; depth++)
	{
		get_child_id(parent_id, &left_id, &right_id);
		NODE* p_left_node   = &kd_tree.p_node[left_id];
		NODE* p_right_node  = &kd_tree.p_node[right_id];
		if (p_pixel[p_node->dim] < p_node->th)
		{
			parent_id = left_id;
			p_node = &kd_tree.p_node[left_id];
		}
		else
		{
			parent_id = right_id;
			p_node = &kd_tree.p_node[right_id];
		}
		index = parent_id;
	}
	if (index < 31)
	{
		int sun = 0;
	}

	return index;
}

void MCUT::init(unsigned char* p_src, int width, int height)
{
	p_src_img  = p_src;
	img_width  = width;
	img_height = height;
}

void MCUT::run_for_image()
{
	create_1d_hist_for_color_image();
	create_kd_tree();
	find_suit_color_for_leaf_nodes_by_src_img();
}



void MCUT::find_suit_color_for_leaf_nodes_by_src_spix(S_Pix* p_spix, int spix_num)
{
	int x = 0, y = 0;
	unsigned char p_pixel[MAX_COLOR_CHANNEL];
	for (int i = 0; i < kd_tree.max_node_num; i++)
	{
		kd_tree.p_node[i].pixel_num = 0;
		for (int j = 0; j < MAX_COLOR_CHANNEL; j++)
		{
			kd_tree.p_node[i].color[j] = 0;
		}
	}


	for (int si = 0; si < spix_num; si++)
	{
		p_pixel[0] = p_spix[si].color.c0;
		p_pixel[1] = p_spix[si].color.c1;
		p_pixel[2] = p_spix[si].color.c2;
		int index = index_kd_tree(p_pixel);
		kd_tree.p_node[index].color[0] += p_pixel[0];
		kd_tree.p_node[index].color[1] += p_pixel[1];
		kd_tree.p_node[index].color[2] += p_pixel[2];
		kd_tree.p_node[index].pixel_num++;
	}

	for (int i = 0; i < kd_tree.max_node_num; i++)
	{
		for (int j = 0; j < MAX_COLOR_CHANNEL && kd_tree.p_node[i].pixel_num > 0; j++)
		{
			kd_tree.p_node[i].color[j] = (int)((float)kd_tree.p_node[i].color[j] / (float)(kd_tree.p_node[i].pixel_num));
		}
	}
}

void MCUT::run_for_superpixel(S_Pix* p_spix, int spix_num)
{
	create_1d_hist_for_spix(p_spix, spix_num);
	create_kd_tree();
	find_suit_color_for_leaf_nodes_by_src_spix(p_spix, spix_num);
}

